"""
基于 AutoGen 的智能体协作流程
"""
import autogen
from autogen import GroupChat, GroupChatManager
from agents.stock_agents import StockAgentFactory, StockAnalysisTeam
from colorama import Fore, Style, init
from typing import Dict, Any, List, Optional
import time
import json
import os


class AutoGenStockFlow:
    """基于 AutoGen 的股票分析智能体协作流程"""
    
    def __init__(self):
        init(autoreset=True)
        print(Fore.CYAN + "🚀 初始化 AutoGen 智能体协作系统...")
        
        # 创建智能体工厂
        self.factory = StockAgentFactory()
        
        # 创建所有智能体
        self.agents = None
        self.team = None
        
        # 加载配置
        self.config = self._load_config()
        
        print(Fore.GREEN + "✅ AutoGen 智能体系统初始化完成")
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置"""
        try:
            config_path = os.path.join(os.path.dirname(__file__), "config.json")
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
                return config.get("autogen_config", {})
        except Exception as e:
            print(Fore.YELLOW + f"⚠️ 加载配置失败，使用默认配置: {e}")
            return {
                "human_input_mode": "NEVER",
                "max_consecutive_auto_reply": 10,
                "clear_history": True
            }
    
    def run_sequential(self, stock_code: str):
        """
        顺序执行模式
        
        Args:
            stock_code: 股票代码
        """
        print(Fore.CYAN + Style.BRIGHT + f"\n🎬 === AutoGen 智能体协作开始（顺序模式）===")
        print(Fore.CYAN + f"📈 目标股票: {stock_code}")
        print(Fore.CYAN + f"🤖 智能体数量: {len(self.agents)} 个")
        print(Fore.CYAN + f"⚡ 处理引擎: AutoGen Sequential")
        
        start_time = time.time()
        
        try:
        self.agents = self.factory.create_all_agents(stock_code)
        self.team = StockAnalysisTeam(self.agents)
            result = self.team.analyze_stock_sequential(stock_code)
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            # 显示结果
            self._display_results(result, processing_time, "顺序执行")
            
            return result
            
        except Exception as e:
            print(Fore.RED + f"\n❌ AutoGen 顺序执行失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def run_group_chat(self, stock_code: str):
        """
        群聊协作模式
        
        Args:
            stock_code: 股票代码
        """
        print(Fore.CYAN + Style.BRIGHT + f"\n🎬 === AutoGen 智能体协作开始（群聊模式）===")
        print(Fore.CYAN + f"📈 目标股票: {stock_code}")
        print(Fore.CYAN + f"🤖 智能体数量: {len(self.agents)} 个")
        print(Fore.CYAN + f"⚡ 处理引擎: AutoGen GroupChat")
        
        start_time = time.time()
        
        try:
        self.agents = self.factory.create_all_agents(stock_code)
        self.team = StockAnalysisTeam(self.agents)
            print(Fore.MAGENTA + f"\n🔄 [AutoGen] 创建智能体群聊...")
            
            # 设置发言顺序
            speaker_transitions = [
                {self.agents["user_proxy"]: [self.agents["data_collector"]]},
                {self.agents["data_collector"]: [self.agents["data_analyst"]]},
                {self.agents["data_analyst"]: [self.agents["decision_maker"]]},
                {self.agents["decision_maker"]: [self.agents["user_proxy"]]},
            ]
            
            # 创建群聊
            groupchat = GroupChat(
                agents=list(self.agents.values()),
                messages=[],
                max_round=15,
                speaker_selection_method="auto",
            )
            
            # 创建群聊管理器
            manager = GroupChatManager(
                groupchat=groupchat,
                llm_config=self.factory.llm_config,
                is_termination_msg=lambda x: "TERMINATE" in x.get("content", ""),
            )
            
            print(Fore.GREEN + f"✅ 群聊创建成功，开始协作分析...")
            
            # 启动群聊
            initial_message = f"""
大家好，我们需要协作分析股票 {stock_code}。

工作流程：
1. 数据收集专家：请收集股票历史数据和相关新闻
2. 数据分析师：基于收集的数据进行深入分析
3. 投资决策专家：根据分析结果制定投资策略

请按顺序开始工作。首先请数据收集专家开始收集 {stock_code} 的数据。
"""
            
            # 用户代理发起对话
            self.agents["user_proxy"].initiate_chat(
                manager,
                message=initial_message,
                clear_history=True
            )
            
            # 获取群聊历史
            chat_history = groupchat.messages
            
            # 生成报告
            result = self._generate_group_chat_report(stock_code, chat_history)
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            # 显示结果
            self._display_results(result, processing_time, "群聊协作")
            
            return result
            
        except Exception as e:
            print(Fore.RED + f"\n❌ AutoGen 群聊协作失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
    
    def run_two_agent_chat(self, stock_code: str):
        """
        两个智能体对话模式（简化版）
        
        Args:
            stock_code: 股票代码
        """
        print(Fore.CYAN + Style.BRIGHT + f"\n🎬 === AutoGen 两智能体对话模式 ===")
        print(Fore.CYAN + f"📈 目标股票: {stock_code}")
        
        start_time = time.time()
        
        try:
        self.agents = self.factory.create_all_agents(stock_code)
        self.team = StockAnalysisTeam(self.agents)
            print(Fore.MAGENTA + "\n🔄 数据收集与分析对话...")
            
            message = f"""
请分析股票 {stock_code}：
1. 收集历史数据和新闻
2. 进行技术和基本面分析
3. 提供投资建议
"""
            
            # 用户代理与数据收集专家对话
            self.agents["user_proxy"].initiate_chat(
                self.agents["data_collector"],
                message=message,
                max_turns=3,
                summary_method="reflection_with_llm"
            )
            
            collection_result = self.agents["user_proxy"].last_message()["content"]
            
            # 数据收集专家与分析师对话
            self.agents["data_collector"].initiate_chat(
                self.agents["data_analyst"],
                message=f"请分析以下数据：\n{collection_result[:1500]}",
                max_turns=2,
                summary_method="reflection_with_llm"
            )
            
            analysis_result = self.agents["data_analyst"].last_message()["content"]
            
            # 分析师与决策专家对话
            self.agents["data_analyst"].initiate_chat(
                self.agents["decision_maker"],
                message=f"基于分析结果，请提供投资决策：\n{analysis_result[:1500]}",
                max_turns=2,
                summary_method="reflection_with_llm"
            )
            
            decision_result = self.agents["decision_maker"].last_message()["content"]
            
            # 生成报告
            result = f"""
{'=' * 60}
股票 {stock_code} 分析报告（两智能体对话模式）
{'=' * 60}

数据收集：
{collection_result[:800]}

数据分析：
{analysis_result[:800]}

投资决策：
{decision_result[:800]}

{'=' * 60}
"""
            
            end_time = time.time()
            processing_time = round(end_time - start_time, 2)
            
            self._display_results(result, processing_time, "两智能体对话")
            return result
            
        except Exception as e:
            print(Fore.RED + f"\n❌ 两智能体对话失败: {str(e)}")
            return None
    
    def _generate_group_chat_report(self, stock_code: str, chat_history: List[Dict]) -> str:
        """生成群聊报告"""
        from datetime import datetime
        
        messages_by_agent = {}
        for msg in chat_history:
            agent_name = msg.get("name")
            if agent_name not in messages_by_agent:
                messages_by_agent[agent_name] = []
            messages_by_agent[agent_name].append(msg.get("content", ""))

        report_sections = []
        for agent_name, messages in messages_by_agent.items():
            section_content = "\n".join(messages)[:1500] if messages else "N/A"
            report_sections.append(f"### {agent_name} ###\n{section_content}")

        report = f"""
        {'=' * 70}
        股票分析报告 - {stock_code}（群聊协作）
        {'=' * 70}

        生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
        分析框架: AutoGen GroupChat
        参与智能体: {len(self.agents)} 个
        对话轮次: {len(chat_history)} 轮

        {"\n\n".join(report_sections)}

        {'=' * 70}
        免责声明：本报告仅供参考，不构成投资建议。
        {'=' * 70}
        """
        return report
    
    def _display_results(self, result: Any, processing_time: float, mode: str):
        """显示处理结果"""
        print(Fore.GREEN + f"\n✅ [AutoGen] {mode}完成 (耗时: {processing_time}s)")
        
        if result:
            print(Fore.WHITE + Style.BRIGHT + "\n📊 === 分析结果 ===")
            print(Fore.WHITE + result[:2500] + "..." if len(result) > 2500 else result)
        
        print(Fore.CYAN + Style.BRIGHT + f"\n🏁 === AutoGen {mode}结束 ===")
    
    def get_agent_info(self):
        """获取智能体信息"""
        if self.agents is None:
            print(Fore.YELLOW + "智能体尚未初始化，将使用默认股票代码 600519.SH 进行初始化。")
            self.agents = self.factory.create_all_agents("600519.SH")
        
        for name, agent in self.agents.items():
            print(f"{Fore.YELLOW}\n{name}:")
            print(f"  - 类型: {type(agent).__name__}")
            print(f"  - 名称: {agent.name}")
            if hasattr(agent, 'system_message'):
                print(f"  - 职责: {agent.system_message[:100]}...")
        
        return self.agents
